import os.path as osp
import xml.etree.ElementTree as ET

from .builder import DATASETS
from .nwpuv2 import NWPUv2Dataset
import mmcv


@DATASETS.register_module()
class NWPUv2DINODataset(NWPUv2Dataset):
    # NWPUv2 Dataset with DINO Attention map

    def __init__(self, **kwargs):
        super(NWPUv2DINODataset, self).__init__(**kwargs)

    def load_annotations(self, ann_file):
        data_infos = []
        img_ids = mmcv.list_from_file(ann_file)
        for img_id in img_ids:
            img_id = img_id.strip()#自己加的
            filename = f'JPEGImages/{img_id}.jpg'
            ftname = f'FeatMap/deg_{img_id}.png'
            xml_path = osp.join(self.img_prefix, 'Annotations',
                                f'{img_id}.xml')
            tree = ET.parse(xml_path)
            root = tree.getroot()
            size = root.find('size')
            width = 0
            height = 0
            if size is not None:
                width = int(size.find('width').text)
                height = int(size.find('height').text)
            else:
                img_path = osp.join(self.img_prefix, 'JPEGImages',
                                    '{}.jpg'.format(img_id))
                img = Image.open(img_path)
                width, height = img.size
                
#             feat_path = osp.join(self.img_prefix, 'FeatMap',
#                                 '{}.png'.format(img_id))
#             feat = Image.open(feat_path)
            
            data_infos.append(
                dict(id=img_id, filename=filename, width=width, height=height, ftname=ftname))
            
        return data_infos